import os
from sklearn.model_selection import train_test_split
import random
import shelve
import tlsh
import copy
from tqdm import tqdm
import hashlib

all_data_path = r"/mnt/mydisk1/chenyongwei/data/functionsim/allorigindata/"
family_lable_Path = r"/home/cyw/projects/function_sim_project/all_data/new_fam_lables.txt"
functionSim_predata = r"/mnt/mydisk1/chenyongwei/sampleDatas/functionSim"
EECG_predata = r"/mnt/mydisk1/chenyongwei/sampleDatas/EECG"

# 保存生成的各个数据集的训练对的信息
sampleLablesPath = "/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/sample_and_lables"
pairInfPath = r"/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/pair_infs"


class EECGdataPre():
    """
        生成训练测试集的样本对
    """

    def __init__(self):
        #   每个样本生成训练队所使用的epoch,epoch越大，数据集越大
        #   为了效率，这里使用都使用1次
        self.train_sample_epoch = 1
        self.test_sample_epoch = 1
        self.valid_sample_epoch = 1

    def get_pair_infs(self, with_Lable=False):
        """
            获取之前生成的的训练对信息
            with_Lable
                True：返回(sample1,sample2,flag,sample1_famlable,sample2_famlable)的数组
                False:返回(sample1,sample2,flag)的数组
        """
        with shelve.open(pairInfPath) as file:
            trainPair = file["trainPair"]
            testPair = file["testPair"]
            validPair = file["validPair"]
        res = {}
        res["train"], res["test"], res["valid"] = [], [], []
        for i in range(len(trainPair)):
            temp = trainPair[i][:3] if with_Lable == False else trainPair
            res["train"].append(temp)
        for i in range(len(testPair)):
            temp = testPair[i][:3] if with_Lable == False else testPair
            res["test"].append(temp)
        for i in range(len(validPair)):
            temp = validPair[i][:3] if with_Lable == False else validPair
            res["valid"].append(temp)
        return res

    def get_all_used_sample(self):
        """
            获得各个模型都能预处理的数据
                能够反汇编 、
                样本原数据存在、
                家族标记存在、
                样本存在startfunc
            返回值：各项数据都存在的数据名称和其家族 res=[],lable=[]
        """
        trans_md5 = self.get_true_md5()
        tempLable = self.get_samples_fam()
        famLable = {}
        modify_Lable_num = 0
        for name in tempLable:
            lable = tempLable[name]
            if name in trans_md5:
                modify_Lable_num += 1
                name = trans_md5[name]
            famLable[name] = lable
        print("修改{}个样本的样本名".format(modify_Lable_num))
        sampleExist = self.get_samples_binary_is_exist()
        disSample = self.get_disassembly_sample()
        res = []
        lable = []
        for i in famLable:
            if i in sampleExist and i in disSample:
                res.append(i)
                lable.append(famLable[i])
        return res, lable

    def samples_train_test_split(self, sampleName, sampleLable):
        """
            获得用于train,test,valid阶段的样本名
            train:test:valid=7:1.5:1.5
        """
        X_train, X_test, y_train, y_test = train_test_split(
            sampleName, sampleLable, test_size=0.3, random_state=666)
        X_valid, X_test, y_valid, y_test = train_test_split(
            X_test, y_test, test_size=0.5, random_state=666)
        print("训练集:{}  测试集：{}  验证集：{}".format(
            len(X_train), len(X_test), len(X_valid)))
        return X_train, y_train, X_test, y_test, X_valid, y_valid

    def generate_sample_pairs(self, name, lable, epoch):
        """
            生成用于训练的训练对，epoch表示循环几轮 
            返回结果为(sample1,sample2,flag,sample1_famlable,sample2_famlable)的数组
            flag=1表示同家族，flag=-1表示不同家族
        """
        Lable_to_sample = {}
        LableList = []
        for i in range(len(lable)):
            if lable[i] not in Lable_to_sample:
                Lable_to_sample[lable[i]] = [name[i]]
                LableList.append(lable[i])
            else:
                Lable_to_sample[lable[i]].append(name[i])
        print("   共有{}种家族".format(len(Lable_to_sample)))
        print("   分别为：{}".format(str(Lable_to_sample.keys())))
        simPair = []
        disSimPair = []
        equelNum = 0
        while epoch != 0:
            epoch -= 1
            for i in range(len(name)):
                family_lable = lable[i]
                size = len(Lable_to_sample[family_lable])
                LableSize = len(LableList)
                # 添加相似的样本对
                if size == 1:
                    # 该家族没有其它的样本，便自己和自己一对
                    simPair.append(
                        (name[i], name[i], 1, family_lable, family_lable))
                    equelNum += 1
                else:
                    # 该家族有其他的样本，找非自己的其他样本组成一对
                    ind = random.randint(0, size-1)
                    tempName = Lable_to_sample[family_lable][ind]
                    while tempName == name[i]:
                        ind = random.randint(0, size-1)
                        tempName = Lable_to_sample[family_lable][ind]
                    simPair.append(
                        (name[i], tempName, 1, family_lable, family_lable))
                # 添加不相似的样本对
                ind = random.randint(0, LableSize-1)
                tempLable = LableList[ind]
                while tempLable == family_lable:
                    ind = random.randint(0, LableSize-1)
                    tempLable = LableList[ind]
                ind = random.randint(0, len(Lable_to_sample[tempLable])-1)
                tempName = Lable_to_sample[tempLable][ind]
                disSimPair.append(
                    (name[i], tempName, -1, family_lable, tempLable))
        res = simPair
        res.extend(disSimPair)
        print("   生成训练对：{}对".format(len(res)))
        print("   相似对中重复：{}对".format(equelNum))
        return res

    def generate_pairs(self):
        """
            划分train，test,valid的样本
            同时生成对应集合种的训练对
        """
        sampleName, sampleLable = self.get_all_used_sample()

        print("各个数据都存在的样本数量为{}".format(len(sampleName)))
        trainName, trainLable, testName, testLable, validName, validLable = self.samples_train_test_split(
            sampleName, sampleLable)

        # 获得各个部分的样本对，按照家族进行分类
        print("生成训练集的样本对")
        trainPair = self.generate_sample_pairs(
            trainName, trainLable, self.train_sample_epoch)
        print("生成测试集的样本对")
        testPair = self.generate_sample_pairs(
            testName, testLable, self.test_sample_epoch)
        print("生成验证集的样本对")
        validPair = self.generate_sample_pairs(
            validName, validLable, self.valid_sample_epoch)

        with shelve.open(sampleLablesPath) as file:
            file["trainName"], file["trainLable"] = trainName, trainLable
            file["testName"], file["testLable"] = testName, testLable
            file["validName"], file["validLable"] = validName, validLable
        with shelve.open(pairInfPath) as file:
            file["trainPair"] = trainPair
            file["testPair"] = testPair
            file["validPair"] = validPair

    def generate_all_to_all_pairs(self, names, lables, saveName):
        """
            生成一个训练集所有样本间的所有匹配对
            及其对应的ground Truth
        """
        #   获得dataInf中的样本和家族信息
        lable_to_name = {}
        name_to_lable = {}
        for i in range(len(names)):
            name_to_lable[names[i]] = lables[i]
        for i in range(len(lables)):
            if lables[i] not in lable_to_name:
                lable_to_name[lables[i]] = [names[i]]
            else:
                lable_to_name[lables[i]].append(names[i])
        #   开始生成目标匹配对
        allPairs = []
        lth = len(names)
        for i in tqdm(range(lth)):
            temp = []
            for j in range(lth):
                if i == j:  # 应该是要提出的,不然mpp值.用哈希的方法会变成1
                    continue
                if lables[i] == lables[j]:
                    temp.append((names[i], names[j], 1))
                else:
                    temp.append((names[i], names[j], -1))
            allPairs.append(copy.deepcopy(temp))
        #   保存数据
        with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/{}".format(saveName)) as file:
            file["pairs"] = allPairs
            file["lable_to_name"] = lable_to_name
            file["name_to_lable"] = name_to_lable

    def get_true_md5(self):
        """
            获得错误md5值和正确md5值之间的映射
        """
        print("正在校验样本md5值与样本名的映射")
        res = {}
        g = os.walk(all_data_path)
        for path, dir_list, file_list in g:
            for file_name in tqdm(file_list):
                # 避免其它中间文件的影响
                if (len(file_name.split(".")) == 1):
                    tar = os.path.join(path, file_name)
                    with open(tar, 'rb') as fp:
                        data = fp.read()
                    file_md5 = hashlib.md5(data).hexdigest()
                    if file_name != file_md5:
                        res[file_md5] = file_name
        return res

    def get_disassembly_sample(self):
        """
            移除结点数大于8000的样本
            将节点数等于0的异常样本移除
            移除radare2不能处理的样本名
            返回剩下的ida处理的样本名
        """
        # ida
        g = os.walk(functionSim_predata)
        cnt = 0
        res = {}
        aNum, bNum = 0, 0
        for path, dir_list, file_list in g:
            for file_name in tqdm(file_list):
                a = file_name.split(".")
                if len(a) == 2 and a[1] == "dir":
                    # 判断该函数是否存在函数结点，不存在结点的函数，保存name，后续查找无结点原因
                    try:
                        with shelve.open(functionSim_predata+"//"+a[0]) as file:
                            cg = file["cg"]
                            if len(cg) == 0:
                                print(
                                    "functionSim反汇编的样本{}---函数结点为{}".format(a[0], len(cg)))
                                aNum += 1
                                continue
                            if len(cg) > 8000:
                                print(
                                    "functionSim反汇编的样本{}---函数结点为{}".format(a[0], len(cg)))
                                bNum += 1
                                continue
                        res[a[0]] = True
                    except Exception as e:
                        print("error:{}".format(e))
        print("ida处理的反汇编样本中,{}个样本不能正常的反汇编,{}个样本过大被移除".format(aNum, bNum))

        # 获得 EECG 能正确预处理的样本
        samples = {}
        for i in os.listdir(EECG_predata):
            tmp = i.split(".")
            if tmp[1] == "dir":
                samples[tmp[0]] = 1
        # 取交集
        ans = {}
        for sample in samples:
            if sample in res:
                ans[sample] = True

        print("反汇编的样本数量为{}".format(len(ans)))
        return ans

    def get_samples_binary_is_exist(self):
        """
            移除掉样本为空输入的情况,使TLSH模型能正常运行
            返回能正常运行的样本名
        """
        print("检验样本是否为空输入")
        res = {}
        g = os.walk(all_data_path)
        for path, dir_list, file_list in g:
            for file_name in tqdm(file_list):
                if tlsh.hash(open(all_data_path+file_name, 'rb').read()) != "TNULL":
                    res[file_name] = True
        print("二进制样本数量为{}".format(len(res)))
        return res

    def get_samples_fam(self):
        """
            获得样本名及其对应的家族
            剔除样本的错误家族：SINGLETON and -	[]
        """
        res = {}
        with open(family_lable_Path) as file:
            for inf in file.readlines():
                temp = inf.split("	")
                name, lable = temp[0], temp[1].strip()
                if name in res:
                    print("样本重复")
                if lable.split(":")[0] == "SINGLETON":
                    print("样本家族标记错误 {} {}".format(name, lable))
                    continue
                if lable.split("	")[0] == "-":
                    print("样本家族标记错误 {} {}".format(name, lable))
                    continue
                res[name] = lable
        print("存在家族的样本数量为{}".format(len(res)))
        return res


if __name__ == "__main__":
    dataProcess = EECGdataPre()
    dataProcess.generate_pairs()
    with shelve.open(sampleLablesPath) as file:
        trainName, trainLable = file["trainName"], file["trainLable"]
        testName, testLable = file["testName"], file["testLable"]
        validName, validLable = file["validName"], file["validLable"]
    print("正在生成all to all pairs")
    dataProcess.generate_all_to_all_pairs(trainName, trainLable, "train")
    dataProcess.generate_all_to_all_pairs(testName, testLable, "test")
    dataProcess.generate_all_to_all_pairs(validName, validLable, "valid")
